import os

from flask import Flask, request, jsonify
from paddleocr import PaddleOCR
from werkzeug.utils import secure_filename

app = Flask(__name__)


@app.route('/')
def hello_world():
    return 'Server Start Successfully!'


def ocr(path):
    """从图片路径获取文字内容"""
    # 初始化PaddleOCR实例，并指定语言为中文（简体）
    ocrC = PaddleOCR(lang='ch')
    try:
        result = ocrC.ocr(path, det=False, cls=False)
        return result
    except Exception as e:
        print(f"错误：{e}")
        return None

@app.route('/ocr', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({ 'code': 400, 'msg': 'No file part.'})
    file = request.files['file']
    if file.filename == '':
        return jsonify({ 'code': 400, 'msg': 'No selected file.'})
    if file:
        # 检查文件扩展名是否为图片格式
        allowed_extensions = {'png', 'jpg', 'jpeg', 'gif'}
        filename = secure_filename(file.filename)
        file_extension = filename.rsplit('.', 1)[1].lower()
        if file_extension not in allowed_extensions:
            return jsonify({ 'code': 400, 'msg': 'Invalid file extension. Only images are allowed.'})

        # 保存上传的文件
        filename = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
        file.save(filename)
        # 调用OCR函数进行识别
        result = ocr(filename)
        if isinstance(result, list):
            # 提取结果中的文字内容，并将其转换为更易于阅读的字符串
            json_data = [[{'text': item[0], 'confidence': item[1]} for item in sublist] for sublist in result]
            # 处理完后删除文件
            os.remove(filename)
            return jsonify({'result': json_data, 'code': 200, 'msg': '识别成功'})
        else:
            return jsonify({ 'code': 500, 'msg': 'OCR识别失败'})


if __name__ == '__main__':
    app.config['UPLOAD_FOLDER'] = 'uploads'
    if not os.path.exists(app.config['UPLOAD_FOLDER']):
        os.makedirs(app.config['UPLOAD_FOLDER'])
    app.run(host='0.0.0.0', port=7985)